#pragma once

#include "Array.h"
#include "Vector.hpp"

/**
 * @file   Matrix.h
 * @author xingyifan <xingyifan@xingyifan-VirtualBox>
 * @date   Wed Jul 14 10:33:43 2021
 *
 * @brief  sparse matrix
 *
 *
 */

class SpMatrix
{
private:
    const int m, n;
    int nnz;
    Array<int> IA, JA;
    Array<double> AA; //对角元放在每行首位

public:
    SpMatrix(int row, int column) : m(row), n(column)
    {
        nnz = std::min(m, n);
        for (int i = 0; i < m; i++)
        {
            IA[i] = i;
        }
        for (int i = 0; i < nnz; i++)
        {
            JA[i] = i;
            AA[i] = 0;
        }
        JA[nnz] = __INT_MAX__;
        IA[m] = nnz;
        AA[nnz] = __FLT_MAX__;
    }
    SpMatrix(int _row, int _column, int _nnz, int *_IA, int *_JA, double *_AA) : m(_row), n(_column)
    {
        nnz = _nnz;
        for (int i = 0; i < m; i++)
        {
            IA[i] = _IA[i];
        }
        IA[m] = nnz;
        for (int i = 0; i < n; i++)
        {
            JA[i] = _JA[i];
        }
        JA[nnz] = __INT_MAX__;
        for (int i = 0; i < nnz; i++)
        {
            AA[i] = _AA[i];
        }
        AA[nnz] = __FLT_MAX__;
    }
    SpMatrix(const SpMatrix &spm) : m(spm.m), n(spm.n)
    {
        nnz = spm.nnz;
        IA = spm.IA;
        JA = spm.JA;
        AA = spm.AA;
    }

// 加减
#define SPMATRIX_BINARY_OP(OpName, Op)                            \
    auto OpName(const SpMatrix &_spm) const                       \
    {                                                             \
        SpMatrix spm = *this;                                     \
        for (int i = 0; i < m; i++)                               \
        {                                                         \
            int range = _spm.IA[i + 1] - _spm.IA[i];              \
            for (int j = 0; j < range; j++)                       \
            {                                                     \
                int col = _spm.IA[i] + j;                         \
                spm.insert(i, _spm.JA[col], Op 1 * _spm.AA[col]); \
            }                                                     \
        }                                                         \
        return spm;                                               \
    }

    SPMATRIX_BINARY_OP(operator+, +)
    SPMATRIX_BINARY_OP(operator-, -)
#undef SPMATRIX_BINARY_OP

    SpMatrix operator*(double k)
    {
        SpMatrix spm = *this;
        for (int i = 0; i < nnz; i++)
        {
            spm.AA[i] *= k;
        }
        return spm;
    }
    //稀疏矩阵的乘法
    SpMatrix operator*(const SpMatrix &_spm)
    {
        const int _m = m;
        const int _n = _spm.n;
        SpMatrix spm(_m, _n);
        for (int i = 0; i < m; i++)
        {
            double tmp_val[_n];
            //计算i行的结果
            for (int j = IA[i]; j < IA[i + 1]; j++)
            {
                //这是一个很巧妙的计算方式，只需要每次计算两个元素的乘积，记录在对应的列的位置上
                //然后不断累加即可，相当于将n次乘法和加法分开计算
                for (int k = _spm.IA[_spm.JA[j]]; k < _spm.IA[_spm.JA[j] + 1]; k++)
                {
                    tmp_val[_spm.JA[k]] += AA[j] * _spm.AA[k];
                }
            }
            //去掉过小的元素
            for (int j = 0; j < _n; j++)
            {
                if (fabs(tmp_val[j]) > __FLT_MIN__)
                {
                    //插入结果
                    spm.insert(i, j, tmp_val[j]);
                }
            }
        }
        return spm;
    }

    friend std::ostream &operator<<(std::ostream &out, const SpMatrix &_spm)
    {
        for (int i = 0; i < _spm.m; i++)
        {
            // IA[m] = nnz，边界值
            for (int j = _spm.IA[i]; j < _spm.IA[i + 1]; j++)
            {
                out << i << ", " << _spm.JA[j] << ", " << _spm.AA[j] << std::endl;
            }
        }
        return out;
    }

    double get(int i, int j) const
    {
        //只能获取存在的位置元素
        for (int k = IA[i]; k < IA[i + 1]; k++)
        {
            if (JA[k] == j)
            {
                return AA[k];
            }
        }
        return 0;
    }

    double &set(int i, int j)
    {
        //只能设置存在的位置
        for (int k = IA[i]; k < IA[i + 1]; k++)
        {
            if (JA[k] == j)
            {
                return AA[k];
            }
        }
        return AA[nnz];
    }

    void insert(int i, int j, double value)
    {
        //如果插入到存在的元素位置，则直接加上插入值
        if (i == j)
        {
            int k = IA[i];
            AA[k] += value;
            return;
        }
        int ins = IA[i] + 1;
        for (int k = ins; k < IA[i + 1]; k++)
        {
            if (JA[k] == j)
            {
                AA[k] += value;
                return;
            }
            else if (JA[k] > j)
            {
                //利用动态数组直接插入
                JA.insert(k, j);
                AA.insert(k, value);
                //改变IA和nnz
                for (int p = i + 1; p < m + 1; p++)
                {
                    IA[p]++;
                }
                nnz++;
                return;
            }
        }
        //当只有对角元的时候插入到下一行开头位置
        JA.insert(ins, j);
        AA.insert(ins, value);
        //改变IA和nnz
        for (int p = i + 1; p < m + 1; p++)
        {
            IA[p]++;
        }
        nnz++;
    }

    SpMatrix transpose()
    {
        SpMatrix res(n, m);
        for (int i = 0; i < m; i++)
        {
            for (int j = frontOfRow(i); j < frontOfRow(i + 1); j++)
            {
                res.insert(posOfCol(j), i, global(j));
            }
        }
        return res;
    }

    int row() const { return m; }
    int col() const { return n; }
    int nonzero() const { return nnz; }

    double diag(int i) { return AA[IA[i]]; }      //获取对角元
    int frontOfRow(int i) const { return IA[i]; } //获取指定行的开头
    int posOfCol(int j) const { return JA[j]; }   //通过全局指标获取元素所在的列位置
    double global(int j) const { return AA[j]; }  //通过全局指标获取元素值

    //获取(i,j)元素的全局列标
    int globalCol(int i, int j) const
    {
        for (int k = IA[i]; k < IA[i + 1]; k++)
        {
            if (JA[k] == j)
            {
                return k;
            }
        }
        return -1;
    }

    // 迭代法
    static Vector<double> Jacobi(SpMatrix &A, Vector<double> &b)
    {
        int n = A.row();
        Vector<double> x(n, 1), y(n);
        double dx;
        do
        {
            dx = 0;
            for (int i = 0; i < n; i++)
            {
                double sum = 0;
                int s_row = A.frontOfRow(i);
                //只需要计算一次加法b[i] / d[i,i]
                sum += b[i] / A.global(s_row);
                for (int j = s_row + 1; j < A.frontOfRow(i + 1); j++)
                {
                    //累计D1(L+U) * x
                    int col = A.posOfCol(j);
                    sum -= A.global(j) * x[col] / A.global(s_row);
                }
                y[i] = sum;
            }
            dx = Vec::norm(y - x, 0);
            x = y;
        } while (dx >= ACCURACY);
        return x;
    }

    static Vector<double> G_S(SpMatrix &A, Vector<double> &b)
    {
        int n = A.row();
        Vector<double> x(n, 1), y(n);
        double dx;
        do
        {
            dx = 0;
            y = x;
            for (int i = 0; i < n; i++)
            {
                double sum = 0;
                int s_row = A.frontOfRow(i);
                //只需要计算一次加法b[i] / d[i,i]
                sum += b[i] / A.global(s_row);
                for (int j = s_row + 1; j < A.frontOfRow(i + 1); j++)
                {
                    //累计D1(L+U) * x
                    int col = A.posOfCol(j);
                    sum -= A.global(j) * x[col] / A.global(s_row);
                }
                // G_S迭代法每次直接更新计算出的x[i]值
                x[i] = sum;
            }
            dx = Vec::norm(y - x, 0);
        } while (dx >= ACCURACY);
        return x;
    }

    static Vector<double> SOR(SpMatrix &A, Vector<double> &b, double w)
    {
        int n = A.row();
        Vector<double> x(n, 1), y(n);
        double dx;
        do
        {
            dx = 0;
            y = x;
            for (int i = 0; i < n; i++)
            {
                double sum = 0;
                int s_row = A.frontOfRow(i);
                //只需要计算一次加法b[i] / d[i,i]
                sum += b[i] / A.global(s_row);
                for (int j = s_row + 1; j < A.frontOfRow(i + 1); j++)
                {
                    //累计D1(L+U) * x
                    int col = A.posOfCol(j);
                    sum -= A.global(j) * x[col] / A.global(s_row);
                }
                // SOR迭代法每次计算出的x[i]值，然后直接乘w修正delta x
                x[i] = sum;
            }
            x = y + (x - y) * w;
            dx = Vec::norm(y - x, 0);
        } while (dx >= ACCURACY);
        return x;
    }
};